'''
Created on Mar 22, 2012

@author: pedro
'''

import Preprocessor as pre
import HDRGenerators
import HDRNormalizers
import sys
import ConfigParser
import os

from Network import OscillatorPyramidalLayer,OutputLayer
from Neurons import PyramidalNeuron,OutputNeuron,integrate_timecourse
from random import shuffle,randint

dt = 0.0002

if __name__ == '__main__':
    
    config_file = sys.argv[1]
    hdr_file = sys.argv[2]
    
    '''---Load Config---'''
    config = ConfigParser.RawConfigParser(allow_no_value=True)
    config.readfp(open(config_file))
    
    time    =    config.getint("heimdallpy", "time")
    dt      =    config.getfloat("heimdallpy", "dt")
    numPyrs =    config.getint("heimdallpy", "numPyrs")
    v0      =    config.getfloat("heimdallpy","v0")
    Rm      =    config.getfloat("heimdallpy","Rm")
    Cm      =    config.getfloat("heimdallpy","Cm")
    tau_ref =    config.getfloat("heimdallpy","tau_ref")
    Vth     =    config.getfloat("heimdallpy","Vth")
    V_spike =    config.getfloat("heimdallpy","V_spike")
    oscillator_freq = config.getfloat("heimdallpy", "oscillator_freq")
    
    stdp_decay_factor   = config.getfloat("output_layer","stdp_decay_factor")
    max_weight          = config.getfloat("output_layer","max_weight")
    connectivity        = config.getfloat("output_layer","connectivity")
    
    Aplus   = config.getfloat("stdp","Aplus")
    Aminus  = config.getfloat("stdp","Aminus")
    
    stdp_tauplus    = config.getfloat("stdp","stdp_tauplus")
    stdp_tauminus   = config.getfloat("stdp", "stdp_tauminus")
    if stdp_tauminus is None:
        stdp_tauminus = stdp_tauplus
    
    time_stop = config.getint("integrate","time_stop")
    '''---End Config---'''
    
    
    '''---Initialize HDR Battery---'''
    
    if os.stat(hdr_file)[6]==0:
        print "Could not find battery in file, creating a new one"
        gen = HDRGenerators.UniformSigmoidHDRGenerator(0.5,5.0)
        #gen = HDRGenerators.RandomizedNVectorHDRGenerator(1.0)
        norm = HDRNormalizers.FloorHDRNormalizer()
        battery = pre.generate_hdr_battery(8, numPyrs, gen, norm)
        pre.save_battery_to_file(battery, hdr_file)
    else:
        print "Loading battery from file...."
        battery = pre.load_battery_from_file(hdr_file)
        print "... loaded battery, found {0} HDRs".format(len(battery))
        
    '''---End Init HDR Battery---'''
    
    
    '''---Initialize Neural Network---'''
    numOuts = numPyrs
    num_trials = 6

    pyramidals = []
    outputs = []
    
    for i in range(numOuts):
        o1 = OutputNeuron(v0,Rm,Cm,tau_ref,Vth,V_spike)
        outputs.append(o1)
        o1 = None
    
    for i in range(numPyrs):
        pyr1 = PyramidalNeuron(v0,Rm,Cm,tau_ref,Vth,V_spike)
        pyramidals.append(pyr1)
        pyr1 = None

    pyrlayer = OscillatorPyramidalLayer(pyramidals, oscillator_freq)
    outlayer = OutputLayer(pyramidals, outputs, stdp_decay_factor, max_weight, connectivity)    
    '''---End Init Neural Network---'''
    
    
    hdr_baseline_outputs = {}
    hdr_pplus_outputs = {}
    
    pplus_hdrs = []
    while len(pplus_hdrs) < 4:
        attempt = randint(1,len(battery)-1)
        if attempt in pplus_hdrs:
            continue
        else:
            pplus_hdrs.append(attempt)
    print "P+ HDRs are : {0}\n".format(pplus_hdrs)
    
    '''---Part 1: Establish Baseline Responses---'''
    
    """Each HDR is presented $num_trials times at random"""
    
    trials_sequence = []
    for i in range(len(battery)):
        for j in range(num_trials):
            trials_sequence.append(i)
    shuffle(trials_sequence)
    
    for hdr_num in trials_sequence:
        #Reset neuron's Vm to V0
        pyrlayer.reset_neurons()
        outlayer.reset_neurons()
        
        #Fire Pyramidal Layer first
        pi_layer_outputs = pyrlayer.fire(battery[hdr_num], dt, time)
        
        #Fire Output Layer Last
        output_layer_outputs = outlayer.fire(pi_layer_outputs, dt, time, True, Aplus, Aminus, stdp_tauplus, stdp_tauminus)
        hdr_baseline_outputs[hdr_num] = sum([integrate_timecourse(output[1], v0, time_stop) for output in output_layer_outputs])
        print "Output for HDR #{0} is {1}".format(hdr_num,hdr_baseline_outputs[hdr_num])    
    
    pyrlayer.reset_neurons()
    outlayer.reset_neurons()
    
    """Acquire Baseline Responses by Firing w/ Learning OFF"""
    for hdr_num in range(len(battery)):
        #Fire Pyramidal Layer first
        pi_layer_outputs = pyrlayer.fire(battery[hdr_num], dt, time)
        
        #Fire Output Layer Last
        output_layer_outputs = outlayer.fire(pi_layer_outputs, dt, time, False, Aplus, Aminus, stdp_tauplus, stdp_tauminus)
        hdr_baseline_outputs[hdr_num] = sum([integrate_timecourse(output[1], v0, time_stop) for output in output_layer_outputs])
        
        pyrlayer.reset_neurons()
        outlayer.reset_neurons()
        
    for hdr_num in range(len(battery)):
        print "Baseline Output for HDR #{0} is {1}".format(hdr_num,hdr_baseline_outputs[hdr_num])
        


    '''---Part 2: Establish P+ Responses---'''
        
        
    for hdr_num in pplus_hdrs:
        
        pyrlayer.reset_neurons()
        outlayer.reset_neurons()
        
        #Fire Pyramidal Layer first
        pi_layer_outputs = pyrlayer.fire(battery[hdr_num], dt, time)
        
        #Fire Output Layer Last
        output_layer_outputs = outlayer.fire(pi_layer_outputs, dt, time, True, Aplus, Aminus, stdp_tauplus, stdp_tauminus)
        hdr_pplus_outputs[hdr_num] = sum([integrate_timecourse(output[1], v0, time_stop) for output in output_layer_outputs])
        print "Output for HDR #{0} is {1}".format(hdr_num,hdr_baseline_outputs[hdr_num])
        
    pyrlayer.reset_neurons()
    outlayer.reset_neurons()
    
    """Acquire End Result Responses by Firing w/ Learning OFF"""
    for hdr_num in range(len(battery)):
        #Fire Pyramidal Layer first
        pi_layer_outputs = pyrlayer.fire(battery[hdr_num], dt, time)
        
        #Fire Output Layer Last
        output_layer_outputs = outlayer.fire(pi_layer_outputs, dt, time, False, Aplus, Aminus, stdp_tauplus, stdp_tauminus)
        hdr_pplus_outputs[hdr_num] = sum([integrate_timecourse(output[1], v0, time_stop) for output in output_layer_outputs])
        
        pyrlayer.reset_neurons()
        outlayer.reset_neurons()
        
    for hdr_num in range(len(battery)):
        print "Baseline Output for HDR #{0} was {1}".format(hdr_num,hdr_baseline_outputs[hdr_num])
        print "Post-PPlus Output for HDR #{0} is {1}".format(hdr_num,hdr_pplus_outputs[hdr_num])
        if hdr_num in pplus_hdrs:
            print "    was a P+ HDR"
    